%% solveKF.m
%
% Solve the KF equation
%--------------------------------------------------------------------------

Mall = Nb*Na*Nz*Nr;

%--------------------------------------------------------------------------
%% Account for Blanchard-Yaari Retirement Process
if deathFlag == 1 & ~exist('DD', 'var')
    %Retire with prob la_BY
    DD = -la_BY*speye(Mall,Mall);
    
    %Switch according to new homeowner distribution, rm = current spot rate
    bMinInd = discretize(0,b);
    bMaxInd = discretize(zstates(1)/2,b);
    wt = la_BY/(bMaxInd-bMinInd+1);
    
    %New homeowners keep same income process (maintains balance) but switch into spot interest rate
    for rm_ind = 1:Nr
        for nz = 1:Nz
            %Finish
            relevant_col_start = M*(rb_spot_ind-1) + Nb*Na*(nz-1) + Nb*(Na-1) + bMinInd;
            relevant_cols = [relevant_col_start:relevant_col_start+bMaxInd-bMinInd];
            
            %Start
            relevant_rows = M*(rm_ind-1) + [1+Nb*Na*(nz-1):Nb*Na*nz];
            relevant_rows = repmat(relevant_rows, length(relevant_cols), 1);
            
            DD = DD + sparse(relevant_rows(:), repmat(relevant_cols',Na*Nb,1), wt, Mall, Mall);
        end
    end
end
if deathFlag == 1
    A = A + DD;    
end


%--------------------------------------------------------------------------
%% Get Intervention Matrices for current rate rb
C = interventionStore{rb_spot_ind};

FF = forcedRefiStore{rb_spot_ind};
A = A + FF;

if max(abs(sum(A, 2))) > 1e-10
    warning('Some rows of A do not sum to 0')
end


%--------------------------------------------------------------------------
%% Combine and solve KF equation
adj = [];
for rm_ind = 1:Nr
    adj = [adj; adjStore{rb_spot_ind,rm_ind}];
end
adj = logical(adj);
adjpoints = pointslistAll(adj);
noadjpoints = pointslistAll(~adj);

if slow_refi ~= 1
    CT = C';
    ACT = transpose(A*C);
    
    % change diagonal of ACT at points in the adjustment region
    AtildeT = ACT + sparse(adjpoints, adjpoints, ones(length(adjpoints),1), Mall, Mall);
elseif slow_refi == 1 && faster_prepay == 0
    %Convert intervention matrix to a slower matrix
    Cslow = C;
    Cslow = Cslow - sparse(noadjpoints, noadjpoints, ones(length(noadjpoints),1), Mall, Mall);
    Cslow = la_slowrefi*Cslow;
    Cslow = Cslow - sparse(adjpoints, adjpoints, la_slowrefi*ones(length(adjpoints),1), Mall, Mall);

    ACT = transpose(A+Cslow);
    AtildeT = ACT; %no difference here anymore
    CT = speye(Mall); %now just keep g as is
elseif slow_refi == 1 && faster_prepay == 1
    %Get prepay/rewrite points
    prepaylist = [];
    for rm_ind = 1:Nr
        tmp = rewrite_combineStore{rb_spot_ind,rm_ind};
        prepaylist = [prepaylist; (1-tmp(:))];
    end
    prepaypoints = pointslistAll(adj&logical(prepaylist));
    rewritepoints = pointslistAll(adj&(~logical(prepaylist)));
   
    %Convert intervention matrix to a slower matrix
    Cslow = C;
    Cslow = Cslow - sparse(noadjpoints, noadjpoints, ones(length(noadjpoints),1), Mall, Mall);
    Cslow(rewritepoints,:) = la_slowrefi*Cslow(rewritepoints,:);
    Cslow(prepaypoints,:) = la_slowprepay*Cslow(prepaypoints,:);
    Cslow = Cslow - sparse(rewritepoints, rewritepoints, la_slowrefi*ones(length(rewritepoints),1), Mall, Mall);
    Cslow = Cslow - sparse(prepaypoints, prepaypoints, la_slowprepay*ones(length(prepaypoints),1), Mall, Mall);

    ACT = transpose(A+Cslow);
    AtildeT = ACT; %no difference here anymore
    CT = speye(Mall); %now just keep g as is
end
    
if findStationary == 1
    if slow_refi ~= 1
        % Fix one value (not in adjustment region) so matrix isn't singular
        vec = zeros(Mall,1);
        rng(10,'twister');
        iFix = noadjpoints(randi(length(noadjpoints),1));
        vec(iFix)=.01;
        AtildeT(iFix,:) = [zeros(1,iFix-1),1,zeros(1,Mall-iFix)];
        
        g_stacked = AtildeT\vec;
    else
        [g_stacked,~] = eigs(AtildeT,1,'lr','SubspaceDimension',25,'MaxIterations',1500);
    end
    g_sum = g_stacked'*ones(Mall,1)*da*db;
    g_stacked = g_stacked./g_sum;
    if max(abs(ACT * g_stacked)) > 1e-6
        warning('Solution to KFE not accurate')
    end
    if sum(g_stacked < -1e-6) ~= 0
        warning('Solution to KFE has negative values')
    end
    g = reshape(g_stacked,Nb,Na,Nz,Nr);

elseif findStationary == 0
    ginit = g_stacked;
    
    if ARMFlag == 1
        %If ARMS, automatically adjust mortgage index
        ARMARM = ARMStore{rb_spot_ind};
        ginit = transpose(ARMARM)*ginit;
    end
    
    g_stacked = bicgstab(speye(Mall)-Delta*ACT, CT*ginit, 1e-6, 750);
    g_sum = g_stacked'*ones(Mall,1)*da*db;
    g_stacked = g_stacked./g_sum;
    if sum(g_stacked < -1e-6) ~= 0
        warning('Solution to KFE has negative values')
    end
    g = reshape(g_stacked,Nb,Na,Nz,Nr);
end
    
